
import json
import time
import string
import random
import base64
from typing import Union, Optional
from datetime import datetime, timedelta

from Crypto.Cipher import AES
from jose import JWTError, jwt
from binascii import b2a_hex, a2b_hex

# JWT算法相关配置
ALGORITHM = "HS512"
ACCESS_TOKEN_EXPIRE_HOURS = 24
SECRET_KEY = "".join(random.choices(f"{string.ascii_letters}{string.digits}", k=64))
ACCESS_TOKEN_KEYS = "".join(random.choices(f"{string.ascii_letters}{string.digits}", k=32))

# AES加密算法相关配置
AES_MODE = AES.MODE_CBC
AES_KEY = "".join(random.choices(f"{string.ascii_letters}{string.digits}", k=16)).encode("utf-8")
AES_IV = "".join(random.choices(f"{string.ascii_letters}{string.digits}", k=16)).encode("utf-8")


def encrypt(data) -> str:
    """
    根据AES算法加密字符串
        :param data: 待加密的数据
    :return: 加密的字符串
    """
    if len(data.encode('utf-8')) % 16:
        add = 16 - (len(data.encode('utf-8')) % 16)
    else:
        add = 0
    text = data + ('\0' * add)
    cryptos = AES.new(AES_KEY, AES_MODE, AES_IV)
    cipher_text = cryptos.encrypt(text.encode('utf-8'))
    return b2a_hex(cipher_text).decode('utf-8')


def decrypt(text) -> str:
    """
    根据AES算法解密字符串
        :param text: 待解密的数据
    :return: 解密的字符串
    """
    cryptos = AES.new(AES_KEY, AES_MODE, AES_IV)
    plain_text = cryptos.decrypt(a2b_hex(text))
    return bytes.decode(plain_text).rstrip('\0')


def generate_encode_token(data: dict, header: dict) -> Optional[str]:
    """
    根据JWT算法加密生成token字符串
        :param data: 待加密的用户数据，包括phone和user_id
        :param header: 追加报文头信息，防止破译
    :return: token字符串
    """
    try:
        if data.get("id", None):
            data["xwi"] = encrypt(str(data["id"]))
            del data["id"]
        if data.get("user_id", None):
            data["xyu"] = encrypt(str(data["user_id"]))
            del data["user_id"]
        if data.get("phone", None):
            data["xxp"] = encrypt(str(data["phone"]))
            del data["phone"]
        expire = datetime.utcnow() + timedelta(hours=ACCESS_TOKEN_EXPIRE_HOURS)
        data.update({"exp": expire, "create_time": time.time()})
        return jwt.encode(data,
                          SECRET_KEY,
                          headers=header,
                          algorithm=ALGORITHM,
                          access_token=ACCESS_TOKEN_KEYS)
    except Exception as err:
        print(err)
        return None


def get_decode_token(x_auth_token: str) -> Optional[dict]:
    """
    根据JWT算法解密token字符串
        :param data: 加密的token
    :return: 用户数据
    """
    try:
        content = jwt.decode(x_auth_token,
                             SECRET_KEY,
                             algorithms=ALGORITHM,
                             access_token=ACCESS_TOKEN_KEYS)
        if content.get("xwi", None):
            content["id"] = decrypt(content["xwi"])
            del content["xwi"]
        if content.get("xyu", None):
            content["user_id"] = decrypt(content["xyu"])
            del content["xyu"]
        if content.get("xxp", None):
            content["phone"] = decrypt(content["xxp"])
            del content["xxp"]
        return content
    except Exception as err:
        print(err)
        return None
